#!coding:utf-8
import numpy as np
import rospy
from sensor_msgs.msg import JointState
from pupper.config import Command, State, BehaviorState, Configuration
from pupper.controller import Controller
from pupper.inverse_kinematics import four_legs_inverse_kinematics
from pupper.joystick_interface import JoystickInterface
from pupper.hardware_interface import HardwareInterface


class PupperControllerNode:
    def __init__(self):
        self.node_name = rospy.get_name()
        rospy.loginfo("[%s] Initializing......" % (self.node_name))
        # 初始化各种对象
        self.config = Configuration()
        self.controller = Controller(
            self.config,
        )
        self.joystick = JoystickInterface(self.config)
        self.state = State()
        command = Command()

        # 初始姿态
        self.state.foot_locations = (
            self.config.default_stance
            + np.array([0, 0, command.height])[:, np.newaxis]
        )
        self.hardware_interface = HardwareInterface()
        # 创建发布者实例，后续可以通过该实例发布对应的消息
        self.joint_states = rospy.Publisher(
            '/joint_states', JointState, queue_size=1)
        # 使用 ros 定时器, 避免手动计算
        rospy.Timer(rospy.Duration(1.0/100), self.step)
        rospy.loginfo("[%s] Initialized." % (self.node_name))

    def step(self, envent=None):
        # 读取手柄数据
        command = self.joystick.get_command(self.state)

        # 执行一次坐标计算
        self.controller.run(self.state, command)

        if self.state.behavior_state != BehaviorState.DEACTIVATED:

            # 坐标到关节角度
            joint_angles = four_legs_inverse_kinematics(
                self.state.rotated_foot_locations, self.config)

            # 关节角度通过JointState话题发布供驱动节点订阅
            self.pub_joint_states(joint_angles)

    def pub_joint_states(self, joint_angles):
        angles = []
        joint_msg = JointState()
        joint_msg.name = ["shoulder_joint_rf", "elbow_joint_rf", "wrist_joint_rf",
                          "shoulder_joint_lf", "elbow_joint_lf", "wrist_joint_lf",
                          "shoulder_joint_rb", "elbow_joint_rb", "wrist_joint_rb",
                          "shoulder_joint_lb", "elbow_joint_lb", "wrist_joint_lb",
                          ]
        for leg_index in range(4):
            for axis_index in range(3):
                angles.append(joint_angles[axis_index][leg_index])
        joint_msg.position = angles
        self.joint_states.publish(joint_msg)

    def onShutdown(self):
        # 节点关闭时要执行的操作
        rospy.loginfo("[%s] Closing ." % (self.node_name))
        rospy.loginfo("[%s] Shutdown." % (self.node_name))


if __name__ == '__main__':
    rospy.init_node('pupper_controller_node')
    node = PupperControllerNode()
    rospy.on_shutdown(node.onShutdown)
    rospy.spin()
